#include "PreRenderProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "Game/Component/MultiModel.h"
#include "LocalScene.h"
#include "LocalEntity.h"

void PreRenderProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mModelEntitys = mLocalScene->GetEntitysByType<LocalEntity, Model>();

	std::vector<Ubo> ubos;
	ubos.resize(mModelEntitys.size());
	for (uint32 i = 0; i < mModelEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mModelEntitys[i]->GetTransform().GetWorldMatrix();

		ubos[i].cameraPosition = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.0);
		ubos[i].lightPosition = glm::vec4(mLocalScene->mLightPosition, 1.0);
		ubos[i].lightColor = glm::vec4(mLocalScene->mLightColor, 1.0);
		ubos[i].lightData.x = mLocalScene->mLightIntensity;
		ubos[i].lightData.y = mLocalScene->mLightRadius;
		ubos[i].lightData.z = static_cast<float>(mModelEntitys[i]->mHasNormal);
	}

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubos](VkCommandBuffer commandBuffer)
		{
			std::vector<Ubo> entityUbos = ubos;

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<PreRenderVertexShader>()
				.PushShader<PreRenderFragmentShader>()
				.SetDescriptorCount(0, mModelEntitys.size())
				.Build();

			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				Model& model = mModelEntitys[i]->GetComponent<Model>();

				std::vector<VertexPosNorUv> tempVertexs;
				tempVertexs.resize(model.mVertexs.size());
				for (uint32 j = 0; j < model.mVertexs.size(); ++j)
				{
					tempVertexs[j].position = model.mVertexs[j].Position;
					tempVertexs[j].normal = model.mVertexs[j].Normal;
					tempVertexs[j].tangent = model.mVertexs[j].Tangent;
					tempVertexs[j].bitangent = model.mVertexs[j].Bitangent;
					tempVertexs[j].uv = model.mVertexs[j].TexC;
				}

				mModelEntitys[i]->mVertexBuffer = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(VertexPosNorUv))
					.Build();

				mModelEntitys[i]->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mIndices.size())
					.Build();

				mModelEntitys[i]->mUboBuffer = CustomBuffer::Builder()
					.SetData(&entityUbos[i])
					.SetSize(sizeof(Ubo))
					.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
					.Build();

				mModelEntitys[i]->mAlbedoImage = Picture2DTexture::Builder()
					.SetPath(mModelEntitys[i]->mAlbedoPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mAlbedoSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mAlbedoImage->GetMipLevels())
					.Build();

				mModelEntitys[i]->mNormalImage = Picture2DTexture::Builder()
					.SetPath(mModelEntitys[i]->mNormalPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mNormalSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mNormalImage->GetMipLevels())
					.Build();

				mModelEntitys[i]->mMetallicImage = Picture2DTexture::Builder()
					.SetPath(mModelEntitys[i]->mMetallicPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mMetallicSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mMetallicImage->GetMipLevels())
					.Build();

				mModelEntitys[i]->mRoughnessImage = Picture2DTexture::Builder()
					.SetPath(mModelEntitys[i]->mRoughnessPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mRoughnessSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mRoughnessImage->GetMipLevels())
					.Build();

				mDescriptor->SetValue("ubo", mModelEntitys[i]->mUboBuffer);
				mDescriptor->SetValue("albedoImage", mModelEntitys[i]->mAlbedoImage, mModelEntitys[i]->mAlbedoSampler);
				mDescriptor->SetValue("normalImage", mModelEntitys[i]->mNormalImage, mModelEntitys[i]->mNormalSampler);
				mDescriptor->SetValue("metallicImage", mModelEntitys[i]->mMetallicImage, mModelEntitys[i]->mMetallicSampler);
				mDescriptor->SetValue("roughnessImage", mModelEntitys[i]->mRoughnessImage, mModelEntitys[i]->mRoughnessSampler);

				mModelEntitys[i]->mDescriptorMask = mDescriptor->AllocateDescriptor();
			}


			mColorImage = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "colorImage");

			mNormalImage = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "normalImage");

			mMetallicImage = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "metallicImage");

			mDepthImage = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "depthImage");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "depthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorImage)
				.AddColorRenderTarget(mNormalImage)
				.AddColorRenderTarget(mMetallicImage)
				.AddColorRenderTarget(mDepthImage)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_FRONT_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorImage)
				.AddRenderTarget(mNormalImage)
				.AddRenderTarget(mMetallicImage)
				.AddRenderTarget(mDepthImage)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "renderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void PreRenderProcess::OnTick(float deltaTime)
{
	std::vector<Ubo> ubos;
	ubos.resize(mModelEntitys.size());
	for (uint32 i = 0; i < mModelEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mModelEntitys[i]->GetTransform().GetLocalMatrix();

		ubos[i].cameraPosition = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.0);
		ubos[i].lightPosition = glm::vec4(mLocalScene->mLightPosition, 1.0);
		ubos[i].lightColor = glm::vec4(mLocalScene->mLightColor, 1.0);
		ubos[i].lightData.x = mLocalScene->mLightIntensity;
		ubos[i].lightData.y = mLocalScene->mLightRadius;
		ubos[i].lightData.z = static_cast<float>(mModelEntitys[i]->mHasNormal);
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubos](VkCommandBuffer commandBuffer)
		{
			std::vector<Ubo> entityUbos = ubos;

			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				mModelEntitys[i]->mUboBuffer->WriteMemory(&entityUbos[i]);
			}


			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (int32 i = 0; i < mModelEntitys.size(); ++i)
			{
				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mModelEntitys[i]->mDescriptorMask.GetDescriptorSet(), 0, nullptr);


				VkBuffer vertexBuffer = mModelEntitys[i]->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mModelEntitys[i]->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdDrawIndexed(commandBuffer, mModelEntitys[i]->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void PreRenderProcess::OnResize()
{
	mColorImage = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "colorImage");

	mNormalImage = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "normalImage");

	mMetallicImage = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "metallicImage");

	mDepthImage = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "depthImage");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
		.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "depthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorImage)
		.AddRenderTarget(mNormalImage)
		.AddRenderTarget(mMetallicImage)
		.AddRenderTarget(mDepthImage)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
